#include <iostream>
#include <new>

using namespace std;

#include <clapack.h>


void SquareBarTorsion(void) {

    // Determine the torsional stress function present in a bar of unit square 
    // cross-section that is fully constrained at one end and subjected to a small twisting moment at the other end. 
    // The twisting moment is applied about an axis perpendicular to the square section.
    // Uses the banded linear solver dgbsv_() from LAPACK.

    // The assembled matrix of FD equations will be held in an array AB.
    // The RHS and solution for the system is held in an array C.
    
    long INFO = -1;  // Holder for the error code returned by dgbsv_()

    const int NI = 8; // Number of nodes in the i-direction.
    const int NJ = 8; // Number of nodes in the j-direction.

    long N = NI*NJ;

    long* IPIV = new long[N];  // Need longs here owing to type requirement for the
                               // LAPACK dgbsv_() function.

    long KL = 8;   // The number of sub-diagonal elements.
    long KU = 8;   // The number of super-diagonal elements.

    long LDAB = 2 * KL + KU + 1;  // Number of rows required by dgbsv_() to represent the 
                                  // banded matrix.

    // Create the banded matrix and initialize it to the default (0.0) using the ().

    double* AB = new double[LDAB * N]();

    // Create right hand side vector and initialize to -1.0/ 128.0.

    long NRHS = 1; // Number of right hand side vectors. One per load case.
    long LDC = N;  // Number of elements in one load case.
    double* C = new double[N * NRHS];

    double cval = -2.0 * (0.5 / (double)(NI)) * (0.5 / (double)(NJ));

    fill_n(C, N * NRHS, cval);

    // Create the 3x8 matrix for the non-zero columns of the B matrix.

    const double B[NJ][3] = {
        0.0, -4.0, 1.0,
        2.0, -4.0, 1.0,
        1.0, -4.0, 1.0,
        1.0, -4.0, 1.0,
        1.0, -4.0, 1.0,
        1.0, -4.0, 1.0,
        1.0, -4.0, 1.0,
        1.0, -4.0, 0.0 };

    // Initizlize AB. Note that the first KL rows of AB should remain zero 
    // for use as workspace by the banded matrix solver dgbsv_().

    // Initialize the first 8 columns of AB.

    int basloc, loc;

    basloc = 0;
    loc = 0;

    for (int j = 0; j < NJ; ++j) {
        loc = basloc + KL;
        AB[loc + 7] = B[j][0];
        AB[loc + 8] = B[j][1];
        AB[loc + 9] = B[j][2];

        loc += 16;

        // The last entry is always 1.0.
        AB[loc] = 1.0;

        basloc += LDAB;
    }

    // Initialize intermediate columns of AB.

    double val_init = 2.0;

    for (int i = 1; i < NI-1; ++i) {

        for (int j = 0; j < NJ; ++j) {

            loc = basloc + KL;
            AB[loc] = val_init;

            AB[loc + 7] = B[j][0];
            AB[loc + 8] = B[j][1];
            AB[loc + 9] = B[j][2];

            loc += 16;

            // The last entry is always 1.0.
            AB[loc] = 1.0;

            basloc += LDAB;
        }

        val_init = 1.0;
    }

    // Initialize the final 8 columns of AB.

    for (int j = 0; j < NJ; ++j) {
        loc = basloc + KL;
        AB[loc] = 1.0;

        AB[loc + 7] = B[j][0];
        AB[loc + 8] = B[j][1];
        AB[loc + 9] = B[j][2];

        basloc += LDAB;
    }

    // Use the LAPACK banded solver to solve [AB]{X} = {C}
    // Note that on exit, the solution {X} is provided in the vector C.

    dgbsv_(&N, &KL, &KU, &NRHS, AB, &LDAB, IPIV, C, &LDC, &INFO);

    cout << "Solution complete. Error code = " << INFO << endl;

    // Tabulate the results.

    cout << '\t' << 'i' << '\t' << 'j' << endl;
    cout << "---------------------" << endl;

    loc = 0;
    for (long i = 0; i < NI; ++i) {
        for (long j = 0; j < NJ; ++j) {
            cout << '\t' << i << '\t' << j << '\t' << C[loc] << endl;
            loc++;
        }
    }

    // Memory clean up.

    delete[] AB;
    delete[] IPIV;
    delete[] C;


}